#include <mpi.h>
#include <iostream>
#include <vector>
#include <cstdlib>
#include <sys/time.h>

using namespace std;

// 随机初始化矩阵
void initializeMatrix(vector<vector<float>>& matrix, int N) {
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
            matrix[i][j] = static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
        }
    }
}

// 执行特殊高斯消去法
void gaussianElimination(int rank, int numProcesses, vector<vector<float>>& matrix, int N) {
    for (int k = 0; k < N; ++k) {
        // 0号进程进行任务分发
        if (rank == 0) {
            // 处理主元行
            for (int j = k + 1; j < N; ++j) {
                matrix[k][j] /= matrix[k][k];
            }
            matrix[k][k] = 1.0;
            // 分发主元行到所有进程
            for (int p = 1; p < numProcesses; ++p) {
                MPI_Send(matrix[k].data(), N, MPI_FLOAT, p, k, MPI_COMM_WORLD);
            }
        } else {
            // 非0号进程接收主元行
            MPI_Recv(matrix[k].data(), N, MPI_FLOAT, 0, k, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }

        // 执行消去
        for (int i = k + 1 + rank; i < N; i += numProcesses) {
            float factor = matrix[i][k];
            for (int j = k + 1; j < N; ++j) {
                matrix[i][j] -= factor * matrix[k][j];
            }
            matrix[i][k] = 0.0;
        }
    }
}

int main(int argc, char* argv[]) {
    MPI_Init(&argc, &argv);

    int numProcesses, rank;
    MPI_Comm_size(MPI_COMM_WORLD, &numProcesses);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    // 设置矩阵大小
    const int N = 1024;

    // 分配矩阵内存
    vector<vector<float>> matrix(N, vector<float>(N));

    // 初始化矩阵
    if (rank == 0) {
        initializeMatrix(matrix, N);
    }

    // 广播矩阵到所有进程
    for (int i = 0; i < N; ++i) {
        MPI_Bcast(matrix[i].data(), N, MPI_FLOAT, 0, MPI_COMM_WORLD);
    }

    // 记录开始时间
    struct timeval t_start, t_end;
    gettimeofday(&t_start, NULL);

    // 执行特殊高斯消去法
    gaussianElimination(rank, numProcesses, matrix, N);

    // 记录结束时间
    gettimeofday(&t_end, NULL);
    if (rank == 0) {
        cout << "Pipeline MPI time cost: "
            << 1000 * (t_end.tv_sec - t_start.tv_sec) +
            0.001 * (t_end.tv_usec - t_start.tv_usec) << "ms" << endl;
    }

    MPI_Finalize();
    return 0;
}
